import os
import open3d as o3d
import yaml
import numpy as np

# 参数
input_path = "../maps/bigmap.pcd"
output_dir = "../maps/multi_pcd/"
written_root_dir = "$(find map_trans_toolkit_pkg)/maps/multi_pcd/"
config_file_path = os.path.join(output_dir, "pcds_config.yaml")
block_size = 10


# 生成配置文件
def create_config_file(output_dir, block_size, root_path, file_names, tfs):
    basic_config = {"root_dir": root_path, "block_size": block_size}

    pcds_config = [
        {"file_name": file_name, "tf": tf}  # [x, y, z, roll, pitch, yaw]
        for file_name, tf in zip(file_names, tfs)
    ]

    config = {"basic": basic_config, "pcds": pcds_config}

    with open(os.path.join(output_dir, "pcds_config.yaml"), "w") as f:
        yaml.dump(config, f)


def split_pcd(input_path, output_dir, block_size, written_root_dir):
    # 创建输出文件夹
    os.makedirs(output_dir, exist_ok=True)

    # 读取点云
    pcd = o3d.io.read_point_cloud(input_path)
    points = np.asarray(pcd.points)

    # 计算点云边界
    min_bound = points.min(axis=0)
    max_bound = points.max(axis=0)

    # 网格划分
    x_bins = np.arange(min_bound[0], max_bound[0], block_size)
    y_bins = np.arange(min_bound[1], max_bound[1], block_size)
    z_bins = np.arange(min_bound[2], max_bound[2], block_size)

    filenames = []
    tfs = []
    count = 0

    for xi in x_bins:
        for yi in y_bins:
            for zi in z_bins:
                # 定义当前立方体的边界
                x_min, x_max = xi, xi + block_size
                y_min, y_max = yi, yi + block_size
                z_min, z_max = zi, zi + block_size

                # 筛选属于当前立方体的点
                mask = (
                    (points[:, 0] >= x_min)
                    & (points[:, 0] < x_max)
                    & (points[:, 1] >= y_min)
                    & (points[:, 1] < y_max)
                    & (points[:, 2] >= z_min)
                    & (points[:, 2] < z_max)
                )

                block_points = points[mask]
                if len(block_points) == 0:
                    continue

                # 保存该块点云
                block_pcd = o3d.geometry.PointCloud()
                block_pcd.points = o3d.utility.Vector3dVector(block_points)

                filename = f"block_{count}.pcd"
                filepath = os.path.join(output_dir, filename)
                o3d.io.write_point_cloud(filepath, block_pcd)

                filenames.append(filename)
                xi,yi = xi+block_size/2, yi+block_size/2
                tfs.append([float(xi), float(yi), float(zi), 0, 0, 0])
                count += 1
                print(f"Saved block {count}: {filename} at position ({xi}, {yi}, {zi})")

    create_config_file(output_dir, block_size, written_root_dir, filenames, tfs)


if __name__ == "__main__":
    split_pcd(input_path, output_dir, block_size, written_root_dir)
